from ultralytics import YOLO
import datetime
import os
import requests
import zipfile

def download_file(url, file_path):
    """下载文件到指定路径"""
    response = requests.get(url, stream=True)
    total_size_in_bytes = int(response.headers.get('content-length', 0))
    progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
    with open(file_path, 'wb') as file:
        for data in response.iter_content(1024):
            progress_bar.update(len(data))
            file.write(data)
    progress_bar.close()

def extract_zip(zip_path, extract_to):
    """解压缩文件到指定目录"""
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(extract_to)

# 初始化YOLO模型，加载预训练权重
model = YOLO('yolov8n.pt')

# 获取当前时间，用于文件名
current_time = datetime.datetime.now().strftime("%Y%m%d%H%M")

# 用户输入显卡型号及数量
gpu_model = input("请输入显卡型号：")
gpu_count = input("请输入显卡数量：")

# 数据集配置文件路径
data_path = 'ultralytics/cfg/datasets/coco128.yaml'
dataset_dir = '/home/amd/datasets/coco128'
zip_path = '/home/amd/datasets/coco128.zip'

# 检查配置文件是否存在，如果不存在则下载并解压数据集
if not os.path.exists(data_path):
    print(f"错误：配置文件 {data_path} 不存在。正在检查数据集...")
    if not os.path.exists(zip_path):
        print(f"{zip_path} 不存在，正在下载...")
        download_file('https://ultralytics.com/assets/coco128.zip', zip_path)
    
    print("正在解压数据集...")
    extract_zip(zip_path, dataset_dir)

    if not os.path.exists(data_path):
        print(f"错误：配置文件 {data_path} 仍然不存在。")
        exit()

epochs = 10  # 训练周期
img_size = 640  # 图像大小
batch_size = 16  # 设置批处理大小
device = '0'  # 使用的GPU设备编号

# 创建一个文件用于保存训练结果
file_name = f"{current_time}.txt"
with open(file_name, 'w') as file:
    file.write("YOLOv8训练用时测试\n")
    file.write(f"显卡型号及数量: {gpu_model} x {gpu_count}\n")

# 进行训练并计算训练用时
start_time = datetime.datetime.now()
results = model.train(data=data_path, batch=batch_size, epochs=epochs, imgsz=img_size, device=device)
end_time = datetime.datetime.now()
training_duration = end_time - start_time

# 将结果写入文件
with open(file_name, 'a') as file:
    file.write(f"\n测试的batch_size: {batch_size}\n")
    file.write(f"实际训练用时: {training_duration}\n")

print(f"批处理大小为 {batch_size} 的训练结果已保存到 {file_name}")